#include "core/memory.h"
#include "tools/log.h"
#include "tools/klib.h"
#include "cpu/mmu.h"
#include "dev/console.h"

static addr_alloc_t paddr_alloc;

static pde_t kernel_page_dir[PDE_CNT] __attribute__((aligned(MEM_PAGE_SIZE)));

static void addr_alloc_init(addr_alloc_t* alloc,uint8_t* bits,uint32_t start,uint32_t size,uint32_t page_size){
    mutex_init(&alloc->mutex);
    alloc->start = start;
    alloc->size = size;
    alloc->page_size = page_size;
    bitmap_init(&alloc->bitmap,bits,alloc->size/page_size,0);
}

//分配多页内存
static uint32_t addr_alloc_page(addr_alloc_t* alloc,int page_count){
    uint32_t addr = 0;
    mutex_lock(&alloc->mutex);

    int page_index = bitmap_alloc_nbits(&alloc->bitmap,0,page_count);
    if(page_index>=0){
        addr = alloc->start+page_index*alloc->page_size;
    }
    mutex_unlock(&alloc->mutex);
    return addr;
}

static void addr_free_page(addr_alloc_t * alloc,uint32_t addr,int page_count){
    mutex_lock(&alloc->mutex);

    uint32_t pg_index =((addr-alloc->start))/alloc->page_size;
    bitmap_set_bit(&alloc->bitmap,pg_index,page_count,0);

    mutex_unlock(&alloc->mutex);
}

void show_mem_info(boot_info_t * boot_info){
    log_printf("mem region:");
    for (int i = 0; i < boot_info->ram_region_count; i++)
    {
        log_printf("[%d]: start at 0x%x -> size 0x%x", i, 
            boot_info->ram_region_cfg[i].start, 
            boot_info->ram_region_cfg[i].size);
    }
    log_printf("\n");
}



//根据vaddr从页表中去查找对应的页表项,alloc表示如果没找到是否新建该项
pte_t* find_pte(pde_t* page_dir,uint32_t vaddr,int alloc){
    pte_t* page_table;
    //page_dir的地址加上当前vaddr中得到的在page_dir中的index,就是vaddr的一级页目录表项
    pde_t* pde = page_dir+pde_index(vaddr);

    //如果表项存在,则直接赋值给page_table
    if(pde->present){
        page_table = (pte_t*)pde_paddr(pde);
    }else{//如果不存在则分配内存,创建表项
        if(alloc==0){
            return (pte_t*)0;
        }
        //如果alloc为1,给还没有的页表分配内存空间用来创建页表
        uint32_t pg_paddr = addr_alloc_page(&paddr_alloc,1);
        if(pg_paddr==0){
            return (pte_t*)0;
        }

        // 设置为用户可读写，将被pte中设置所覆盖
        pde->v = pg_paddr | PDE_P | PDE_W | PDE_U ;

        // 为物理页表绑定虚拟地址的映射，这样下面就可以计算出虚拟地址了
        //kernel_pg_last[pde_index(vaddr)].v = pg_paddr | PTE_P | PTE_W;

        // 清空页表，防止出现异常
        // 这里虚拟地址和物理地址一一映射，所以直接写入
        page_table = (pte_t*)pg_paddr;
        kernel_memset(page_table,0,MEM_PAGE_SIZE);
    }
    //一级页目录表项的地址加上vaddr中的二级表项的index,得到二级页表中的表项pte的地址
    return page_table + pte_index(vaddr);
}

/**
 * @brief 将指定的地址空间进行一页的映射
 */
int memory_create_map(pde_t* page_dir,uint32_t vaddr,uint32_t paddr,int count,uint32_t perm){
    for (int i = 0; i < count; i++)
    {
        pte_t* pte = find_pte(page_dir,vaddr,1);
        if(pte==(pte_t*)0){
            return -1;
        }
        ASSERT(pte->present==0);
        //设置页表项pte的权限
        pte->v = paddr|perm|PTE_P;

        vaddr+=MEM_PAGE_SIZE;
        paddr+=MEM_PAGE_SIZE;
    }
    return 0;
}
//创建内核页表
void create_kernel_table(){
    extern uint8_t* kernel_base[],s_text[],e_text[],s_data[];

    // 地址映射表, 用于建立内核级的地址映射
    // 地址不变，但是添加了属性
    static memory_map_t kernel_map[] = {
        {kernel_base,s_text,kernel_base,PTE_W},     // 内核栈区
        {s_text,e_text,s_text,0},           // 内核代码区
        {s_data,(void*)(MEM_EBDA_START-1),s_data,PTE_W},     // 内核数据区
        {(void*)CONSOLE_DISP_ADDR,(void*)CONSOLE_DISP_END,(void*)CONSOLE_DISP_ADDR,PTE_W},
        {(void*)MEME_EXT_START,(void*)MEM_EXT_END,(void*)MEME_EXT_START,PTE_W}
    };

    // 清空后，然后依次根据映射关系创建映射表
    //也就是内核在内存中的代码段，数据段等在内存中所占的位置，然后将其的虚拟地址和物理地址的对应关系映射好
    //然后再建立页表，建立分页映射关系
    for (int i = 0; i < sizeof(kernel_map)/sizeof(memory_map_t); i++)
    {
        memory_map_t * map = kernel_map + i;

        uint32_t vstart = down2((uint32_t)map->vstart,MEM_PAGE_SIZE);
        uint32_t vend = up2((uint32_t)map->vend,MEM_PAGE_SIZE);

        int page_count = (vend-vstart)/MEM_PAGE_SIZE;

        //将内核使用的栈区,代码区和数据区的内存空间建立到页表的映射关系
        memory_create_map(kernel_page_dir,vstart,(uint32_t)map->pstart,page_count,map->perm);
    }
    
}

static uint32_t total_mem_size(boot_info_t* boot_info){
    uint32_t mem_size = 0;
    for (int i = 0; i < boot_info->ram_region_count; i++)
    {
        mem_size+=boot_info->ram_region_cfg[i].size;
    }
    return mem_size;
}

uint32_t memory_create_uvm(){
    pde_t* page_dir = (pde_t*)addr_alloc_page(&paddr_alloc,1);
    if(page_dir==0){
        return 0;
    }
    kernel_memset((void*)page_dir,0,MEM_PAGE_SIZE);
    //得到用户进程起始位置在pde表中的索引
    uint32_t user_pde_start = pde_index(MEMORY_TASK_BASE);
    //映射内核地址
    for(int i=0;i<user_pde_start;i++){
        page_dir[i].v = kernel_page_dir[i].v;
    }
    return (uint32_t)page_dir;
}

void memory_init(boot_info_t * boot_info){
    extern uint8_t* mem_free_start;

    uint8_t* mem_free = (uint8_t*)&mem_free_start;

    log_printf("memory init...");
    show_mem_info(boot_info);
    
    //内存1M以上的内存大小,1M一下不使用,作为操作系统使用
    uint32_t mem_up1MB_free = total_mem_size(boot_info) - MEME_EXT_START;
    //将可用内存按页大小取整
    mem_up1MB_free = down2(mem_up1MB_free,MEM_PAGE_SIZE);

    log_printf("free memory: start at: 0x%x, size: 0x%x", MEME_EXT_START, mem_up1MB_free);

    // 4GB大小需要总共4*1024*1024*1024/4096/8=128KB的位图, 使用低1MB的RAM空间中足够
    // 该部分的内存仅跟在mem_free_start开始放置,因此这里的mem_free实际上就是内核内存空间后的内存位置，
    //该位置也在内存的不到1M的空间处，用mem_free作为存放位图数据的bits数组
    addr_alloc_init(&paddr_alloc,mem_free,MEME_EXT_START,mem_up1MB_free,MEM_PAGE_SIZE);

    //总的可用内存按页分,一共的个数,每个占一个bit位,这里计算出位图占的内存大小,再加在位图的开始地址上
    //让mem_free的地址跳过位图所占的内存空间,此时mem_free的位置应该在bitmap后,MEM_EBDA_START之前
    mem_free += bitmap_byte_count(paddr_alloc.size/MEM_PAGE_SIZE);

    ASSERT(mem_free < (uint8_t *)MEM_EBDA_START);

    create_kernel_table();

    mmu_set_page_dir((uint32_t)kernel_page_dir);
}

int memory_alloc_for_page_dir(uint32_t page_dir,uint32_t vaddr, uint32_t size,int perm){
    //记录下当前的地址分配到哪里了
    uint32_t curr_vaddr = vaddr;    
    //计算出一共要分配的页数
    int page_count = up2(size,MEM_PAGE_SIZE) / MEM_PAGE_SIZE;

    for (int i = 0; i < page_count; i++)
    {
        //分配一页物理内存
        uint32_t paddr = addr_alloc_page(&paddr_alloc,1);
        if(paddr==0){
            log_printf("mem page alloc failed,no memory");
            return 0 ;
        }
        //建立物理内存和虚拟内存的映射关系,也就是将物理页的地址写到页表中的pte项中
        int err = memory_create_map((pde_t*)page_dir,curr_vaddr,paddr,1,perm);
        if(err<0){
            log_printf("create memory failed. err=%d",err);
            return -1;
        }

        curr_vaddr += MEM_PAGE_SIZE;
    }
    
    return 0;
}

int memory_alloc_page_for(uint32_t addr,uint32_t size,int perm){
    //使用具体的哪个pagedir来分配页
    return memory_alloc_for_page_dir(task_current()->tss.cr3,addr,size,perm);
}

uint32_t memory_alloc_page(){
    uint32_t addr = addr_alloc_page(&paddr_alloc,1);
    return addr;
}

static pde_t* curr_page_dir(){
    return (pde_t*)(task_current()->tss.cr3);
}
void memory_free_page(uint32_t addr){
    if(addr<MEMORY_TASK_BASE){
        addr_free_page(&paddr_alloc,addr,1);
    }else{
        pte_t* pte = find_pte(curr_page_dir(),addr,0);
        ASSERT((pte == (pte_t*)0)&&pte->present);

        addr_free_page(&paddr_alloc,pte_paddr(pte),1);
    }
}

/**
 * @brief 销毁用户空间内存
 */
void memory_destory_uvm (uint32_t page_dir) {
    uint32_t user_pde_start = pde_index(MEMORY_TASK_BASE);
    pde_t * pde = (pde_t *)page_dir + user_pde_start;

    ASSERT(page_dir != 0);

    // 释放页表中对应的各项，不包含映射的内核页面
    for (int i = user_pde_start; i < PDE_CNT; i++, pde++) {
        if (!pde->present) {
            continue;
        }

        // 释放页表对应的物理页 + 页表
        pte_t * pte = (pte_t *)pde_paddr(pde);
        for (int j = 0; j < PTE_CNT; j++, pte++) {
            if (!pte->present) {
                continue;
            }

            addr_free_page(&paddr_alloc, pte_paddr(pte), 1);
        }

        addr_free_page(&paddr_alloc, (uint32_t)pde_paddr(pde), 1);
    }

    // 页目录表
    addr_free_page(&paddr_alloc, page_dir, 1);
}

/**
 * @brief 复制页表及其所有的内存空间
 */
uint32_t memory_copy_uvm (uint32_t page_dir) {
    // 复制基础页表
    uint32_t to_page_dir = memory_create_uvm();
    if (to_page_dir == 0) {
        goto copy_uvm_failed;
    }

    // 再复制用户空间的各项
    uint32_t user_pde_start = pde_index(MEMORY_TASK_BASE);
    pde_t * pde = (pde_t *)page_dir + user_pde_start;

    // 遍历用户空间页目录项
    for (int i = user_pde_start; i < PDE_CNT; i++, pde++) {
        if (!pde->present) {
            continue;
        }

        // 遍历页表
        pte_t * pte = (pte_t *)pde_paddr(pde);
        for (int j = 0; j < PTE_CNT; j++, pte++) {
            if (!pte->present) {
                continue;
            }

            // 分配物理内存
            uint32_t page = addr_alloc_page(&paddr_alloc, 1);
            if (page == 0) {
                goto copy_uvm_failed;
            }

            // 建立映射关系
            uint32_t vaddr = (i << 22) | (j << 12);
            int err = memory_create_map((pde_t *)to_page_dir, vaddr, page, 1, get_pte_perm(pte));
            if (err < 0) {
                goto copy_uvm_failed;
            }

            // 复制内容。
            kernel_memcpy((void *)page, (void *)vaddr, MEM_PAGE_SIZE);
        }
    }
    return to_page_dir;

copy_uvm_failed:
    if (to_page_dir) {
        memory_destory_uvm(to_page_dir);
    }
    return -1;
}


uint32_t memory_get_paddr (uint32_t page_dir, uint32_t vaddr) {
    pte_t * pte = find_pte((pde_t *)page_dir, vaddr, 0);
    if (pte == (pte_t *)0) {
        return 0;
    }

    return pte_paddr(pte) + (vaddr & (MEM_PAGE_SIZE - 1));
}

// uint32_t memory_get_paddr(uint32_t page_dir,uint32_t vaddr){
//     pte_t* pte = find_pte((pde_t*)page_dir,vaddr,0);
//     if(!pte){
//         return 0;
//     }
//     return pte_paddr(pte) + (vaddr & (MEM_PAGE_SIZE-1));
// }


int memory_copy_uvm_data(uint32_t to,uint32_t page_dir,uint32_t from,uint32_t size){
    while (size > 0)
    {
        uint32_t to_paddr = memory_get_paddr(page_dir,to);
        if(to_paddr == 0){
            return -1;
        }

        uint32_t offset_in_page = to_paddr & (MEM_PAGE_SIZE-1);
        uint32_t curr_size = MEM_PAGE_SIZE - offset_in_page;
        if(curr_size>size){
            curr_size = size;
        }

        kernel_memcpy((void*)to_paddr,(void*)from,curr_size);

        size -= curr_size;
        to += curr_size;
        from +=curr_size;
    }
    return 0;
}

char * sys_sbrk(int incr) {
    task_t * task = task_current();
    char * pre_heap_end = (char * )task->heap_end;
    int pre_incr = incr;

    ASSERT(incr >= 0);

    // 如果地址为0，则返回有效的heap区域的顶端
    if (incr == 0) {
        //log_printf("sbrk(0): end = 0x%x", pre_heap_end);
        return pre_heap_end;
    } 
    
    uint32_t start = task->heap_end;
    uint32_t end = start + incr;

    // 起始偏移非0
    int start_offset = start % MEM_PAGE_SIZE;
    if (start_offset) {
        // 不超过1页，只调整
        if (start_offset + incr <= MEM_PAGE_SIZE) {
            task->heap_end = end;
            return pre_heap_end;
        } else {
            // 超过1页，先只调本页的
            uint32_t curr_size = MEM_PAGE_SIZE - start_offset;
            start += curr_size;
            incr -= curr_size;
        }
    }

    // 处理其余的，起始对齐的页边界的
    if (incr) {
        uint32_t curr_size = end - start;
        int err = memory_alloc_page_for(start, curr_size, PTE_P | PTE_U | PTE_W);
        if (err < 0) {
            log_printf("sbrk: alloc mem failed.");
            return (char *)-1;
        }
    }

    //log_printf("sbrk(%d): end = 0x%x", pre_incr, end);
    task->heap_end = end;
    return (char * )pre_heap_end;        
}